import pika
import json
from functools import partial

DESERIALIZATION_FUNCTIONS = {
    'json': json.loads
}


class RabbitMQBroker(object):

    def __init__(self, username=None, password=None, host=None, logger=None, deserialization_method=None):
        self.host = host
        self.username = username
        self.password = password

        self.conn = None
        self.channel = None
        self.logger = logger
        self.deserialization_method = deserialization_method

        self.setup()

    def setup(self):
        if self.username and self.password:
            credentials = pika.PlainCredentials(username=self.username, password=self.password)
            connection_params = pika.ConnectionParameters(host=self.host, credentials=credentials)
        else:
            connection_params = pika.ConnectionParameters(host=self.host)

        self.conn = pika.BlockingConnection(connection_params)
        self.channel = self.conn.channel()
        self.channel.basic_qos(prefetch_count=1)

    def consume(self, queues, callback):
        for queue in queues:
            self.channel.queue_declare(queue=queue, durable=True)
            callback = partial(callback, queue=queue)
            self.channel.basic_consume(callback=callback, queue=queue)

    def publish(self, exchange, key, message):
        self.channel.basic_publish(exchange, key, json.dumps(message, ensure_ascii=False))
        self.logger.info("Publish message {} {}  {}".format(exchange, key, message))

    def basic_ack(self, delivery_tag):
        self.channel.basic_ack(delivery_tag=delivery_tag)

    def close(self):
        if self.conn:
            self.conn.close()

    def start_consuming(self):
        return self.channel.start_consuming

    def deserialize_body(self, body):
        if not self.deserialization_method:
            return body

        deserialization_func = DESERIALIZATION_FUNCTIONS[self.deserialization_method]

        try:
            body = deserialization_func(body)
        except Exception as err:
            self.logger.info("{} deserialization body {} error {}".format(self.__class__.__name__, body, err))
        return body


def get_broker_from_kwargs(logger, **kwargs):
    host = kwargs.get('host')
    username = kwargs.get('username')
    password = kwargs.get('password')
    queue_sensor_config = kwargs.get('rabbitmq_queue_sensor')

    queues = queue_sensor_config['queues']
    if not isinstance(queues, list):
        queues = [queues]
    deserialization_method = queue_sensor_config['deserialization_method']
    return RabbitMQBroker(username=username,
                          password=password,
                          host=host,
                          logger=logger,
                          deserialization_method=deserialization_method), queues
